import os 
import platform
import torch
import torch.nn as nn
from collections import OrderedDict
from torchvision import models
from ._resnet import (resnet18, resnet34, resnet50, resnet101, resnet152, resnext50_32x4d, resnext101_32x8d, wide_resnet50_2, wide_resnet101_2)
from ._senet import se_resnext50, se_resnet50, se_resnet101, se_resnext101


class MobileNet_V2(nn.Module):
    def __init__(self, num_classes, pretrained=None):
        super(MobileNet_V2, self).__init__()
        model = models.mobilenet_v2(pretrained=pretrained)
        model.classifier = nn.Sequential(
                nn.Dropout(p=0.2, inplace=False),
                nn.Linear(1280, num_classes, bias=True),
        )
        self.net = model

    def forward(self, img):
        output = self.net(img)
        return output


class SeResNext_50(nn.Module):
    def __init__(self, input_channels=3, num_classes=2, pretrained=None):
        super(SeResNext_50, self).__init__()
        model = se_resnext50(num_classes, pretrained=pretrained)
        if input_channels != 3:
            model.layer0.conv1 = nn.Conv2d(input_channels, 64,
                              kernel_size=7, stride=2,
                              padding=3, bias=False)
        self.net = model

    def forward(self, img):
        _, output = self.net(img)
        return output

    # def extract(self, img):
    #     feature_map, _ = self.net(img)
    #     return feature_map

class SeResNext_101(nn.Module):
    def __init__(self, input_channels=3, num_classes=2, pretrained=None):
        super(SeResNext_101, self).__init__()
        model = se_resnext101(num_classes, pretrained=pretrained)
        if input_channels != 3:
            model.layer0.conv1 = nn.Conv2d(input_channels, 64,
                              kernel_size=7, stride=2,
                              padding=3, bias=False)
        self.net = model

    def forward(self, img):
        _, output = self.net(img)
        return output

    # def extract(self, img):
    #     feature_map, _ = self.net(img)
    #     return feature_map


class SeResNet_50(nn.Module):
    def __init__(self, input_channels=3, num_classes=2, pretrained=None):
        super(SeResNet_50, self).__init__()
        model = se_resnet50(num_classes, pretrained=pretrained)
        if input_channels != 3:
            model.layer0.conv1 = nn.Conv2d(input_channels, 64,
                            kernel_size=7, stride=2,
                            padding=3, bias=False)
        self.net = model

    def forward(self, img):
        output = self.net(img)
        return output


class ResNet_50(nn.Module):
    def __init__(self, input_channels=3, num_classes=2, ):
        super(ResNet_50, self).__init__()
        model = resnet50(pretrained=True)
        if input_channels != 3:
            model.conv1 = nn.Conv2d(input_channels, 64,
                              kernel_size=7, stride=2,
                              padding=3, bias=False)
        model.fc = nn.Sequential(
                nn.Linear(2048, num_classes, bias=True),
        )
        self.net = model

    def forward(self, img):
        output = self.net(img)
        return output


class ResNext_50(nn.Module):
    def __init__(self, input_channels=3, num_classes=2, pretrained=False):
        super(ResNext_50, self).__init__()
        model = resnext50_32x4d(pretrained=pretrained)
        if input_channels != 3:
            model.conv1 = nn.Conv2d(input_channels, 64,
                            kernel_size=7, stride=2,
                            padding=3, bias=False)
        if num_classes != 1000:
            model.fc = nn.Sequential(
                    nn.Linear(2048, num_classes, bias=True),
            )
        self.net = model
    
    def forward(self, img):
        output = self.net(img)
        return output


class ResNet_18(nn.Module):
    def __init__(self, input_channels=3, num_classes=2, pretrained=False):
        super(ResNet_18, self).__init__()
        pretrained = True if (pretrained is not None) else False
        model = resnet18(pretrained=True)
        # model.conv1 = nn.Conv2d(input_channels, 64,
        #                   kernel_size=7, stride=2,
        #                   padding=3, bias=False)
        model.fc = nn.Sequential(
                nn.Linear(512, num_classes, bias=True),
        )
        self.net = model

    def forward(self, img):
        output = self.net(img)
        return output